import tensorflow as tf
from tensorflow.python.saved_model import saved_model, signature_constants, tag_constants
#from tensorflow.lite.python import convert_saved_model
from tensorflow_core.lite.python import convert_saved_model

import os

def create_saimple_savedModel():
    saved_model_dir = os.path.join('data', 'simple_savedmodel')
    with tf.Session() as sess:
        in_tensor = tf.placeholder(dtype=tf.float32, shape=[2,2])
        w = tf.Variable(initial_value=[[1,2],[3,4]], dtype=tf.float32)
        out_tensor = in_tensor + w
        inputs = {'x':in_tensor}
        outputs = {'y':out_tensor}
        sess.run(tf.global_variables_initializer())
        saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
    return saved_model_dir

def convert_savedmodel2pb(saved_model_dir, input_arrays=None, input_shapes=None, output_arrays=None, tag_set=None, signature_key=None):
    if tag_set is None:
        tag_set = set([tag_constants.SERVING])
    if signature_key is None:
        signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    graph_def, in_tensors, out_tensors, _ = (convert_saved_model.freeze_saved_model(saved_model_dir=saved_model_dir,
            input_arrays=input_arrays,
            input_shapes=input_shapes,
            output_arrays=output_arrays,
            tag_set=tag_set,
            signature_key=signature_key))
    print(graph_def)
    print(in_tensors)
    print(out_tensors)
    return graph_def, in_tensors, out_tensors

def save_pb(grap_def):
    tf.train.write_graph(grap_def, 'data/simple_pb', 'pb.pbtxt', as_text=True)


if __name__ == '__main__':
    #create_saimple_savedModel()
    graph_def, _, _ = convert_savedmodel2pb(os.path.join('data', 'simple_savedmodel'))
    save_pb(graph_def)